import numpy as np
import matplotlib.pyplot as plt
import sounddevice as sd
from scipy.fft import fft, ifft, fftfreq
from scipy.io.wavfile import write
from pydub import AudioSegment

# --- 1. 参数设置 ---
SAMPLE_RATE = 44100  # 采样率 (Hz)
DURATION = 5  # 录音时长 (秒)
CUTOFF_FREQ = 1000  # 截止频率 (Hz)，您可以随时调整这个值来进行实验


def fourier_transform_demo():
    """
    一个完整的傅里叶变换音频处理演示函数 (VBR版本, 绘图功能已修复)
    """
    # --- 2. 录制音频 (带倒计时) ---
    print(f"准备录音，请在倒数结束后开始说话或唱歌...")
    for i in range(3, 0, -1):
        print(f"{i}...")
        sd.sleep(1000)
    print("开始录音！")

    myrecording = sd.rec(int(DURATION * SAMPLE_RATE), samplerate=SAMPLE_RATE, channels=1, dtype='float32')
    sd.wait()
    print("录音结束。")

    original_signal = myrecording.flatten()

    # --- 3. 计算傅里叶变换 ---
    N = len(original_signal)
    yf = fft(original_signal)
    xf = fftfreq(N, 1 / SAMPLE_RATE)

    # --- 4. 高频滤波 ---
    yf_filtered = yf.copy()
    yf_filtered[np.abs(xf) > CUTOFF_FREQ] = 0

    # --- 5. 逆傅里叶变换 ---
    filtered_signal = ifft(yf_filtered).real

    # --- 6. 可视化 (完整的绘图部分) ---
    print("正在生成可视化图表...")
    fig, axs = plt.subplots(4, 1, figsize=(10, 12))
    fig.suptitle('傅里叶变换音频滤波演示', fontsize=16)

    time_axis = np.linspace(0, DURATION, N)

    # 子图1: 原始时域信号
    axs[0].plot(time_axis, original_signal, lw=1)
    axs[0].set_title("1. 原始音频波形 (时域)")
    axs[0].set_xlabel("时间 (s)")
    axs[0].set_ylabel("振幅")
    axs[0].grid(True)

    # 子图2: 原始频域信号 (频谱)
    positive_freq_mask = xf >= 0
    axs[1].plot(xf[positive_freq_mask], np.abs(yf[positive_freq_mask]))
    axs[1].set_title("2. 原始音频频谱 (频域)")
    axs[1].set_xlabel("频率 (Hz)")
    axs[1].set_ylabel("能量")
    axs[1].set_xlim(0, SAMPLE_RATE / 4)
    axs[1].grid(True)

    # 子图3: 滤波后的频域信号
    axs[2].plot(xf[positive_freq_mask], np.abs(yf_filtered[positive_freq_mask]))
    axs[2].set_title(f"3. 滤掉 {CUTOFF_FREQ} Hz 以上高频后的频谱")
    axs[2].set_xlabel("频率 (Hz)")
    axs[2].set_ylabel("能量")
    axs[2].axvline(x=CUTOFF_FREQ, color='r', linestyle='--', label=f'截止频率: {CUTOFF_FREQ} Hz')
    axs[2].legend()
    axs[2].set_xlim(0, SAMPLE_RATE / 4)
    axs[2].grid(True)

    # 子图4: 滤波后的时域信号
    axs[3].plot(time_axis, filtered_signal, lw=1)
    axs[3].set_title("4. 滤波后还原的音频波形 (时域)")
    axs[3].set_xlabel("时间 (s)")
    axs[3].set_ylabel("振幅")
    axs[3].grid(True)

    plt.tight_layout(rect=[0, 0.03, 1, 0.96])
    plt.show()

    # --- 7. 播放对比 ---
    print("\n准备播放原始音频...")
    sd.play(original_signal, SAMPLE_RATE)
    sd.wait()

    sd.sleep(1000)

    print(f"\n准备播放滤掉 {CUTOFF_FREQ} Hz 以上高频后的音频...")
    sd.play(filtered_signal, SAMPLE_RATE)
    sd.wait()

    # --- 8. 保存文件 (使用VBR) ---
    print("\n正在将音频保存到文件...")

    original_signal_int16 = np.int16(original_signal / np.max(np.abs(original_signal)) * 32767)
    filtered_signal_int16 = np.int16(filtered_signal / np.max(np.abs(filtered_signal)) * 32767)

    write("original_audio.wav", SAMPLE_RATE, original_signal_int16)
    write("filtered_audio.wav", SAMPLE_RATE, filtered_signal_int16)

    audio_segment = AudioSegment(
        filtered_signal_int16.tobytes(),
        frame_rate=SAMPLE_RATE,
        sample_width=2,
        channels=1
    )
    audio_segment.export("filtered_audio.mp3", format="mp3", parameters=['-q:a', '2'])

    print("所有文件已保存！")


if __name__ == '__main__':
    fourier_transform_demo()